package websocket

import (
	"encoding/json"
	"fmt"
	"net/http"
	"sync"
	"time"

	"github.com/gorilla/websocket"
	"network-firewall-go/internal/types"
)

// Server WebSocket服务器
type Server struct {
	clients    map[*Client]bool
	register   chan *Client
	unregister chan *Client
	broadcast  chan []byte
	upgrader   websocket.Upgrader
	mutex      sync.RWMutex
}

// Client WebSocket客户端
type Client struct {
	conn   *websocket.Conn
	send   chan []byte
	server *Server
	id     string
}

// NewServer 创建新的WebSocket服务器
func NewServer() *Server {
	return &Server{
		clients:    make(map[*Client]bool),
		register:   make(chan *Client),
		unregister: make(chan *Client),
		broadcast:  make(chan []byte),
		upgrader: websocket.Upgrader{
			CheckOrigin: func(r *http.Request) bool {
				// 允许所有来源（生产环境中应该更严格）
				return true
			},
			ReadBufferSize:  1024,
			WriteBufferSize: 1024,
		},
	}
}

// Start 启动WebSocket服务器
func (s *Server) Start() {
	go s.run()
	fmt.Println("✅ WebSocket服务器已启动")
}

// run 运行WebSocket服务器主循环
func (s *Server) run() {
	for {
		select {
		case client := <-s.register:
			s.mutex.Lock()
			s.clients[client] = true
			s.mutex.Unlock()
			
			fmt.Printf("🔗 WebSocket客户端连接: %s\n", client.id)
			
			// 发送欢迎消息
			welcome := &types.WebSocketMessage{
				Type: "welcome",
				Data: map[string]interface{}{
					"message":   "连接成功",
					"client_id": client.id,
					"timestamp": time.Now().Unix(),
				},
			}
			s.sendToClient(client, welcome)

		case client := <-s.unregister:
			s.mutex.Lock()
			if _, ok := s.clients[client]; ok {
				delete(s.clients, client)
				close(client.send)
			}
			s.mutex.Unlock()
			
			fmt.Printf("🔌 WebSocket客户端断开: %s\n", client.id)

		case message := <-s.broadcast:
			s.mutex.RLock()
			for client := range s.clients {
				select {
				case client.send <- message:
				default:
					delete(s.clients, client)
					close(client.send)
				}
			}
			s.mutex.RUnlock()
		}
	}
}

// HandleWebSocket 处理WebSocket连接
func (s *Server) HandleWebSocket(w http.ResponseWriter, r *http.Request) {
	conn, err := s.upgrader.Upgrade(w, r, nil)
	if err != nil {
		fmt.Printf("❌ WebSocket升级失败: %v\n", err)
		return
	}

	client := &Client{
		conn:   conn,
		send:   make(chan []byte, 256),
		server: s,
		id:     generateClientID(),
	}

	s.register <- client

	// 启动客户端处理goroutine
	go client.writePump()
	go client.readPump()
}

// BroadcastTransaction 广播事务
func (s *Server) BroadcastTransaction(transaction *types.HTTPTransaction) {
	message := &types.WebSocketMessage{
		Type: "transaction",
		Data: transaction,
	}
	
	s.broadcastMessage(message)
}

// BroadcastStatus 广播状态更新
func (s *Server) BroadcastStatus(status *types.ProxyStatus) {
	message := &types.WebSocketMessage{
		Type: "status",
		Data: status,
	}
	
	s.broadcastMessage(message)
}

// BroadcastRuleUpdate 广播规则更新
func (s *Server) BroadcastRuleUpdate(rule *types.Rule) {
	message := &types.WebSocketMessage{
		Type: "rule_update",
		Data: rule,
	}
	
	s.broadcastMessage(message)
}

// broadcastMessage 广播消息
func (s *Server) broadcastMessage(message *types.WebSocketMessage) {
	data, err := json.Marshal(message)
	if err != nil {
		fmt.Printf("❌ 序列化WebSocket消息失败: %v\n", err)
		return
	}

	select {
	case s.broadcast <- data:
	default:
		fmt.Println("⚠️ WebSocket广播通道已满，丢弃消息")
	}
}

// sendToClient 发送消息给特定客户端
func (s *Server) sendToClient(client *Client, message *types.WebSocketMessage) {
	data, err := json.Marshal(message)
	if err != nil {
		fmt.Printf("❌ 序列化WebSocket消息失败: %v\n", err)
		return
	}

	select {
	case client.send <- data:
	default:
		s.mutex.Lock()
		delete(s.clients, client)
		close(client.send)
		s.mutex.Unlock()
	}
}

// GetClientCount 获取连接的客户端数量
func (s *Server) GetClientCount() int {
	s.mutex.RLock()
	defer s.mutex.RUnlock()
	return len(s.clients)
}

// readPump 读取客户端消息
func (c *Client) readPump() {
	defer func() {
		c.server.unregister <- c
		c.conn.Close()
	}()

	// 设置读取超时
	c.conn.SetReadDeadline(time.Now().Add(60 * time.Second))
	c.conn.SetPongHandler(func(string) error {
		c.conn.SetReadDeadline(time.Now().Add(60 * time.Second))
		return nil
	})

	for {
		messageType, message, err := c.conn.ReadMessage()
		if err != nil {
			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
				fmt.Printf("❌ WebSocket读取错误: %v\n", err)
			}
			break
		}

		if messageType == websocket.TextMessage {
			c.handleMessage(message)
		}
	}
}

// writePump 写入消息到客户端
func (c *Client) writePump() {
	ticker := time.NewTicker(54 * time.Second)
	defer func() {
		ticker.Stop()
		c.conn.Close()
	}()

	for {
		select {
		case message, ok := <-c.send:
			c.conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
			if !ok {
				c.conn.WriteMessage(websocket.CloseMessage, []byte{})
				return
			}

			w, err := c.conn.NextWriter(websocket.TextMessage)
			if err != nil {
				return
			}
			w.Write(message)

			// 批量发送队列中的其他消息
			n := len(c.send)
			for i := 0; i < n; i++ {
				w.Write([]byte{'\n'})
				w.Write(<-c.send)
			}

			if err := w.Close(); err != nil {
				return
			}

		case <-ticker.C:
			c.conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
			if err := c.conn.WriteMessage(websocket.PingMessage, nil); err != nil {
				return
			}
		}
	}
}

// handleMessage 处理客户端消息
func (c *Client) handleMessage(message []byte) {
	var msg types.WebSocketMessage
	if err := json.Unmarshal(message, &msg); err != nil {
		fmt.Printf("❌ 解析WebSocket消息失败: %v\n", err)
		return
	}

	switch msg.Type {
	case "ping":
		// 响应ping消息
		pong := &types.WebSocketMessage{
			Type: "pong",
			Data: map[string]interface{}{
				"timestamp": time.Now().Unix(),
			},
		}
		c.server.sendToClient(c, pong)

	case "subscribe":
		// 处理订阅请求
		fmt.Printf("📡 客户端 %s 订阅: %v\n", c.id, msg.Data)

	case "unsubscribe":
		// 处理取消订阅请求
		fmt.Printf("📡 客户端 %s 取消订阅: %v\n", c.id, msg.Data)

	default:
		fmt.Printf("⚠️ 未知WebSocket消息类型: %s\n", msg.Type)
	}
}

// generateClientID 生成客户端ID
func generateClientID() string {
	return fmt.Sprintf("client_%d", time.Now().UnixNano())
} 